#!/usr/bin/env python3                                                                                                                                                                                          
# encoding: utf-8
import os
import pickle 
from multiprocessing.connection import Listener
from multiprocessing.connection import Client
from threading import Thread
 
from fabric.colors import magenta, green, yellow, cyan, red
from config import HADOOP_USER
 
 
class RPCHandler:
    def __init__(self):
        self._functions = {}
 
    def register_function(self, func):
        self._functions[func.__name__] = func
 
    def handle_connection(self, connection): 
        try:
            while True:
        # Receive a message
                func_name, args, kwargs = pickle.loads(connection.recv()) 
                # Run the RPC and send a response
                try:
                    r = self._functions[func_name](*args,**kwargs)
                    connection.send(pickle.dumps(r))
                except Exception as e:
                    connection.send(pickle.dumps(e)) 
        except EOFError:
            pass
 
 
class RPCProxy:
    
    def __init__(self, connection):
        self._connection = connection
 
    def __getattr__(self, name):
        def do_rpc(*args, **kargs):
            self._connection.send(pickle.dumps((name, args, kargs)))
            result = pickle.loads(self._connection.recv())
            if isinstance(result, Exception):
                raise result
            return result
        return do_rpc
 
 
def get_client(address, auth_token):
    r = RPCProxy(Client(address, authkey=auth_token))
    r.re_address = address
    return r
 
 
def interact(client):
         
    def copy(local, remote):
        if os.path.exists(local):
            remote_path = os.path.join(remote, local.split("/")[-1])
            with open(local, 'rb') as fp:
                client.copyToRemote(remote_path, fp.read())
         
    def check_path(path):
        ps = path.split("/")
        if ps[-1] == ".":
            return "/".join(ps[:-1])
        elif ps[-1] == "..":
            return "/".join(ps[:-2])
        else:
            return path
         
    def check_cmd(remote_pwd, cmd):
        args = cmd.split()
        if cmd.startswith("inlocal"):
            local_args = args[1:]
            print(green(os.popen(" ".join(local_args)).read()))
            return True
         
        if cmd.startswith("hdfs"):
            hdfs_arg = ' '.join(args[1:])
            stat, msg = client.hdfs_cmd(hdfs_arg, prefix="cd %s && " % remote_pwd, user=HADOOP_USER, path=remote_pwd)
            if stat != 2:
                print(cyan("%s" % (hdfs_arg)))
            return True
         
        if cmd.startswith("put"):
            local_files = args[1:]
            remote_dir = remote_pwd
            for l in local_files:
                copy(l, remote_dir)
                print(cyan("upload: %s => %s" % (l, remote_dir)))
            return True
        if cmd.startswith("cd"):
            return False
         
        stat, res =  client.Bash("cd %s && %s " % (remote_pwd, cmd))
        if stat == 0:
            print(yellow(res))
        elif stat == 2:
            print(red("[err]: %s " % res)) 
        else:
            print(res)
         
        return False
            
         
    import re
    print(green("this host : %s , Ctrl-c to exit\n" % client.re_address[0]))
    print(green("you can use such special prefix cmd\n" ))
    print("""
inlocal : will excute in local for the other cmd ,which been in same line . example "inlocal ls ./"
put : upload local files to remote's path , which is terminated by context  . example "put f1 f2 f3 f4  "
toHdfs : pass file in remote to remote's hdfs file system """)

    remote_pwd = "/"
    while True:
        try:
            cmd = input(magenta("%s >> " % remote_pwd))
            if not cmd.startswith("inlocal"):
                remote_pwd = os.path.join(remote_pwd, re.findall(r'(?:cd\s+?)([\w_\/\-\.]+)', cmd)[0]) if cmd.find("cd") != -1 else remote_pwd
            else:
                pass
            remote_pwd = check_path(remote_pwd)
            check_cmd(remote_pwd, cmd)
        except KeyboardInterrupt as e:
            print(green("bye bye!!"))
            break
        except EOFError as e:
            print(green("bye bye!!"))
            break
                                                                                                                                                                                                                
        
        
def rpc_server(handler, address, authkey): 
    sock = Listener(address, authkey=authkey) 
    while True:
        client = sock.accept()
        t = Thread(target=handler.handle_connection, args=(client,))
        t.daemon = True
        t.start()
